package com.zr.flink;

import com.example.klinemqflink.config.RabbitMQConfig;
import com.example.klinemqflink.model.MarketData;
import com.example.klinemqflink.model.RMDeserializationSchema;
import org.apache.flink.api.common.eventtime.TimestampAssignerSupplier;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.OpenContext;
import org.apache.flink.configuration.StateBackendOptions;
import org.apache.flink.contrib.streaming.state.RocksDBOptions;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.PrintSinkFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.connectors.rabbitmq.RMQSource;
import org.apache.flink.streaming.connectors.rabbitmq.common.RMQConnectionConfig;
import org.apache.flink.util.Collector;
import com.example.klinemqflink.model.Kline;

import java.time.Duration;
import java.util.List;

public class KLineCalculationJob {

    public static void main(String[] args) throws Exception {
        // 1. 启用 RocksDB 状态后端（减少内存压力）
        Configuration config = new Configuration();
        config.setString(StateBackendOptions.STATE_BACKEND, "rocksdb");
        config.setString(RocksDBOptions.LOCAL_DIRECTORIES, "/Users/zhangrui/flink_cache"); // 指定本地存储路径

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config);
        env.setParallelism(1); // 根据CPU核心数调整
        env.enableCheckpointing(10); // 启用检查点

        RMQConnectionConfig connectionConfig = new RMQConnectionConfig.Builder()
                .setHost("localhost")
                .setPort(5672)
                .setUserName("guest")
                .setPassword("guest")
                .setVirtualHost("/")
                .build();
        // 定义多个队列名称
        List<String> queueNames = RabbitMQConfig.MARKET_DATA_QUEUES;
        DataStream<MarketData> combinedStream = null;

        // 为每个队列创建 RMQSource 并合并数据流
        for (String queueName : queueNames) {
            RMQSource<MarketData> source = new RMQSource<>(
                    connectionConfig,
                    queueName,
                    new RMDeserializationSchema()
            );

            DataStream<MarketData> stream = env.addSource(source);

            if (combinedStream == null) {
                combinedStream = stream;
            } else {
                combinedStream = combinedStream.union(stream);
            }
        }

        // 3. 定义时间窗口（事件时间）
        WatermarkStrategy<MarketData> watermarkStrategy = WatermarkStrategy.<MarketData>forBoundedOutOfOrderness(Duration.ofSeconds(5)).withTimestampAssigner((event, timestamp) -> event.getTimestamp());

        combinedStream.assignTimestampsAndWatermarks(watermarkStrategy).keyBy(MarketData::getSymbol) // 按交易对分组
                .window(TumblingEventTimeWindows.of(Time.minutes(1))) // 1分钟窗口
                .aggregate(new IncrementalKLineAggregate(), new KLineWindowFunction()).name("1m KLine").addSink(new PrintSinkFunction<>());

        combinedStream.assignTimestampsAndWatermarks(watermarkStrategy).keyBy(MarketData::getSymbol) // 按交易对分组
                .window(TumblingEventTimeWindows.of(Time.minutes(5))) // 1分钟窗口
                .aggregate(new IncrementalKLineAggregate(), new KLineWindowFunction()).name("5m KLine").addSink(new PrintSinkFunction<>());

        env.execute("Real-time K-Line Calculation");
    }

    // 增量聚合函数（优化内存）
    public static class IncrementalKLineAggregate implements AggregateFunction<MarketData, KLineAccumulator, KLineAccumulator> {

        @Override
        public KLineAccumulator createAccumulator() {
            return new KLineAccumulator();
        }

        @Override
        public KLineAccumulator add(MarketData value, KLineAccumulator accumulator) {
            if (accumulator.getOpen() == 0) {
                accumulator.setOpen(value.getPrice());
                accumulator.setLow(value.getPrice());
                accumulator.setHigh(value.getPrice());
            } else {
                accumulator.setLow(Math.min(accumulator.getLow(), value.getPrice()));
                accumulator.setHigh(Math.max(accumulator.getHigh(), value.getPrice()));
            }
            accumulator.setClose(value.getPrice());
            accumulator.setVolume(accumulator.getVolume() + value.getVolume());
            accumulator.setCount(accumulator.getCount() + 1);
            return accumulator;
        }

        @Override
        public KLineAccumulator merge(KLineAccumulator a, KLineAccumulator b) {
            a.setLow(Math.min(a.getLow(), b.getLow()));
            a.setHigh(Math.max(a.getHigh(), b.getHigh()));
            a.setClose(b.getClose());
            a.setVolume(a.getVolume() + b.getVolume());
            a.setCount(a.getCount() + b.getCount());
            return a;
        }

        @Override
        public KLineAccumulator getResult(KLineAccumulator accumulator) {
            return accumulator;
        }
    }

    // 窗口处理函数（生成完整K线）
    public static class KLineWindowFunction extends ProcessWindowFunction<KLineAccumulator, Kline, String, TimeWindow> {



        @Override
        public void process(String symbol, ProcessWindowFunction<KLineAccumulator, Kline, String, TimeWindow>.Context context, Iterable<KLineAccumulator> iterable, Collector<Kline> collector) throws Exception {
            KLineAccumulator acc = iterable.iterator().next();
            Kline kline = new Kline(symbol, context.window().getStart(), context.window().getEnd(), acc.getOpen(), acc.getHigh(), acc.getLow(), acc.getClose(), acc.getVolume());
            collector.collect(kline);
        }

        @Override
        public void open(OpenContext openContext) throws Exception {
            super.open(openContext);
        }
    }

    // K线累加器
    public static class KLineAccumulator {
        private double open;
        private double high;
        private double low;
        private double close;
        private double volume;
        private int count;

        public double getOpen() {
            return open;
        }

        public void setOpen(double open) {
            this.open = open;
        }

        public double getHigh() {
            return high;
        }

        public void setHigh(double high) {
            this.high = high;
        }

        public double getLow() {
            return low;
        }

        public void setLow(double low) {
            this.low = low;
        }

        public double getClose() {
            return close;
        }

        public void setClose(double close) {
            this.close = close;
        }

        public double getVolume() {
            return volume;
        }

        public void setVolume(double volume) {
            this.volume = volume;
        }

        public int getCount() {
            return count;
        }

        public void setCount(int count) {
            this.count = count;
        }
    }

}